########################################################################
#  Copyright (C) 2007,8,9 Ehud Ben-Reuven
#  udi@benreuven.com
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation version 2.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
########################################################################
"""Convert OFX v1.x or QFX files to flat lines.
This is done in two stages.
The first converts the OFX to N3 file that keeps the structure of the original
OFX file.
The second stage converts the N3 file to a graph and then unstructure it.
"""
help="""Convert OFX v1.x file (MS-Money or Quicken Web Connect) to
flat OFX-RDF (in memory) and load to Sesame server
"""
import logging
import ofx2n3
from rdflib.graph import *
from rdflib.namespace import *
from rdflib.term import *
OFX = Namespace("http://www.w3.org/2000/10/swap/pim/ofx#")
A3 = Namespace("http://code.google.com/p/3account/wiki/Schema#")
XSD = Namespace("http://www.w3.org/2001/XMLSchema#")

FS = Namespace("file://")
flatLine = A3["flatLine"]
NS={'ofx':OFX, 'a3':A3}

def getSimplifiedName(uri):
    if "#" in uri: return uri[uri.rfind("#") + 1:]
    return uri[uri.rfind("/") + 1:]

def pre(fname,fout,hashing=False):
    ofx2n3.flat(fname,fout,hashing=hashing)

def cleanval(v):
    if not isinstance(v,URIRef):
        return v
    v=str(v)
    for ns in NS.itervalues():
        if v.startswith(str(ns)):
            return Literal(v[len(str(ns)):])
    return Literal(v)

def cleanvar(l):
    l = l[l.rfind(":")+1:] if ":" in l else l
    l = l if l[0].isalpha() else l[1:]
    return l

def vars(lst):
    res=[]
    for l in lst:
        assert isinstance(l,str)
        res.append("?"+cleanvar(l))
    return " ".join(res)

def pairs(x,lst):
    res=[]
    for l in lst:
        assert isinstance(l,str)
        if l[0]=="+":
            l=l[1:]
            res[-2]+=" ; ofx:"+l.upper()+" ?"+l
        else:
            if l[0]=="?":
                sstr=" optional { ?"+x+" "
                estr=" } .\n"
                l=l[1:]
            else:
                sstr="?"+x+" "
                estr=" .\n"
            while ":" in l:
                sstr+="ofx:"+l[:l.find(":")]+" [ "
                estr=" ]"+estr
                l=l[l.find(":")+1:]
            sstr+="ofx:"+l.upper()+" ?"+l
            res.append(sstr)
            res.append(estr)
    return " ".join(res)
    
def bprt(f,bNode,lst,vals):
    for i,t in enumerate(lst):
        v=vals[i]
        if v:
            t=cleanvar(t)
            v=cleanval(v)
            if t.isupper():
                t=OFX[t]
            else:
                t=A3[t]
            f.add((bNode,t,v))
            
def flatseclist(g,f,context):
    secinfoT=["SECID:UNIQUEID","+UNIQUEIDTYPE","SECNAME",]
    secinfoO=["?TICKER","?FIID","?RATING","?UNITPRICE","?dtasof",
              "?CURRENCY:CURRATE","+CURSYM","?MEMO"]
    infoTs = [("DEBT",["PARVALUE",
                       "DEBTTYPE",
                       "?DEBTCLASS",
                       "?COUPONRT",
                       "?DTCOUPON","?COUPONFREQ",
                       "?CALLPRICE","?YIELDTOCALL","?DTCALL","?CALLTYPE",
                       "?YIELDTOMAT","?DTMAT","?ASSETCLASS","?FIASSETCLASS"
                     ]),
            ("MF",["?MFTYPE","?YIELD","?DTYIELDASOF",
                   "?MFASSETCLASS:PORTION:ASSETCLASS","+PERCENT",
                   "?FIMFASSETCLASS:FIPORTION:FIASSETCLASS","+PERCENT"]),
            ("OPT",["OPTTYPE","STRIKEPRICE","DTEXPIRE","SHPERCTRCT",
                    "?secid","?ASSETCLASS","?FIASSETCLASS"]),
            ("OTHER",["?TYPEDESC","?ASSETCLASS","?FIASSETCLASS"]),
            ("STOCK",["?STOCKTYPE","?YIELD","?DTYIELDASOF","?ASSETCLASS",
                      "?FIASSETCLASS"])]
    for sectype,infoT in infoTs:
        secinfos=g.query("select "+vars(secinfoT+secinfoO+infoT)+" {\n"+
                    "[ ofx:SECLIST [ ?SECTYPE ?x ] ] .\n"+
                    "?x ofx:SECINFO ?y .\n"+
                    pairs("y",secinfoT)+
                    pairs("x",infoT)+
                    pairs("y",secinfoO)+ # crazy bug in librdf 2.4.2
                    "}",
                    initNs=NS,initBindings={"SECTYPE":OFX[sectype+"INFO"]})
        for secinfo in secinfos:
            bNode=BNode()
            f.add((context,flatLine,bNode))
            f.add((bNode,
                   OFX["SECTYPE"],
                   Literal(getSimplifiedName(sectype))))
            bprt(f,bNode,secinfoT+secinfoO+infoT,secinfo)

def flatinv(g,f,acctid,context):
    invstmtrT=["INVACCTFROM:ACCTID","+BROKERID","CURDEF","DTASOF"]
    invstmtrs=g.query("select "+vars(invstmtrT)+"{"+pairs("x",invstmtrT)+"}",
                      initNs=NS,initBindings={"ACCTID":acctid})
    assert len(invstmtrs)==1
    for invstmtr in invstmtrs:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        f.add((bNode,A3["aCCTTYPE"],Literal("INVESTMENT")))
    #INVPOSLIST
    invposT=["SECID:UNIQUEID","+UNIQUEIDTYPE",
             "HELDINACCT","POSTYPE","UNITS","UNITPRICE","MKTVAL","DTPRICEASOF",
             "?CURRENCY:CURRATE","+CURSYM","?memo"]
    invposs=g.query("select ?pos "+vars(invposT)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    "?x ofx:INVPOSLIST [ ?pos [ ofx:INVPOS ?y ] ] . "+
                    pairs("y",invposT)+"}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invpos in invposs:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,["pos"]+invposT,invpos)
    #INVBAL
    invbalT=["INVBAL:AVAILCASH","+MARGINBALANCE","+SHORTBALANCE"]
    invbals=g.query("select "+vars(invbalT)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    pairs("x",invbalT)+"}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invbal in invbals:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,invbalT,invbal)
    #INVTRANLIST
    invtranlistT=["DTSTART","DTEND"]
    invtranlists=g.query("select "+vars(invtranlistT)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    "?x ofx:INVTRANLIST ?y . "+
                    pairs("y",invtranlistT)+"}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invtranlist in invtranlists: # At most one INVTRANLIST per account
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,invtranlistT,invtranlist)
    
    #INVTRANLIST:INVBUY/SELL
    invtranT=["INVTRAN:FITID","+DTTRADE","SECID:UNIQUEID","+UNIQUEIDTYPE"]
    invtranTopt=[
        "?INVTRAN:SRVRTID","?INVTRAN:DTSETTLE","?INVTRAN:REVERSALFITID",
        "?INVTRAN:mEMO",];
    invbuysellT=invtranT+[
             "units","UNITPRICE",
             "TOTAL",
             "SUBACCTSEC","SUBACCTFUND",
             ]
    invbuysellTopt=invtranTopt+[
             "?COMMISSION","?TAXES","?FEES","?LOAD",
             "?CURRENCY:CURRATE","+CURSYM",
             "?ORIGCURRENCY:currate","+cursym",
             "?INV401KSOURCE","?LOANID",
             "?WITHHOLDING","?TAXEXEMPT","?GAIN",# INVSELL
             "?STATEWITHHOLDING","?PENALTY","?MARKDOWN",
             "?MARKUP", # INVBUY
             "?LOANPRINCIPAL","?LOANINTEREST",
             "?DTPAYROLL","?PRIORYEARCONTRIB",
             ]
    invtranOpTopt=["?ACCRDINT","?BUYTYPE","?RELFITID","?OPTBUYTYPE","?SHPERCTRCT",
              "?SELLREASON","?SELLTYPE","?AVGCOSTBASIS","?OPTSELLTYPE",
              "?RELTYPE","?SECURED",]
    invbuysells=g.query("select ?tran ?tranOp "+
                        vars(invbuysellT)+vars(invbuysellTopt)+
                        vars(invtranOpTopt)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    "?x ofx:INVTRANLIST [ ?tran ?z ] ."+
                    "?z ?tranOp ?y ."+
                    pairs("y",invbuysellT)+
                    pairs("y",invbuysellTopt)+
                    pairs("z",invtranOpTopt)+
                    "}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invbuysell in invbuysells:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,["tran","tranOp"]+invbuysellT+invbuysellTopt+invtranOpTopt,
             invbuysell)
    #INVTRANLIST:INCOME
    incomeT=invtranT
    incomeTopt=invtranTopt+["?OPTACTION","?units","?SHPERCTRCT","?SUBACCTSEC",
                      "?RELFITID","?GAIN","?INCOMETYPE","?TOTAL",
                      "?SUBACCTFUND","?TAXEXEMPT","?WITHHOLDING",
                      "?CURRENCY:CURRATE","+CURSYM",
                      "?ORIGCURRENCY:CURRATE","+CURSYM",
                      "?SUBACCTTO","?SUBACCTFROM",
                      "?UNITPRICE","?COMMISSION","?TAXES","?FEES","?LOAD",
                      "?OLDUNITS","?NEWUNITS","?NUMERATOR","?DENOMINATOR",
                      "?FRACCASH","?TFERACTION","?POSTYPE",
                      "?INVACCTFROM:acctid","+brokerid", # lower case not to coflict with the transaction's ACCTID,BROKERID
                      "?AVGCOSTBASIS","?UNITPRICE","?DTPURCHASE"
                      ]
    incomes=g.query("select ?tran "+
                        vars(incomeT)+vars(incomeTopt)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    "?x ofx:INVTRANLIST [ ?tran ?y ] ."+
                    pairs("y",incomeT)+
                    pairs("y",incomeTopt)+
                    "}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for income in incomes:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,["tran"]+incomeT+incomeTopt,income)
    #INVTRANLIST:INVBANKTRAN
    invbanktranT=["SUBACCTFUND"]
    stmttrnT=["TRNTYPE","DTPOSTED","TRNAMT","FITID",
              "?DTUSER","?DTAVAIL",
              "?CORRECTFITID","?CORRECTACTION","?SRVRTID",
              "?CHECKNUM","?REFNUM","?SIC","?PAYEEID","?NAME",
              "?PAYEE:NAME","+ADDR1","+ADDR2","+ADDR3","+CITY","+STATE","+POSTALCODE","+COUNTRY","+PHONE",
              "?BANKACCTTO:bankid","+acctid","+accttype",
              "?CCACCTTO:acctid",
              "?mEMO","?INV401KSOURCE",
              "?CURRENCY:currate","+cursym",
              "?ORIGCURRENCY:currate","+cursym",
              ]
    invbanktrans=g.query("select "+vars(invbanktranT)+vars(stmttrnT)+"{"+
                    pairs("x",["INVACCTFROM:ACCTID"])+
                    "?x ofx:INVTRANLIST [ ofx:INVBANKTRAN ?y ] ."+
                    "?y ofx:STMTTRN ?z ."+
                    pairs("y",invbanktranT)+pairs("z",stmttrnT)+
                    "}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invbanktran in invbanktrans:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,invbanktranT+stmttrnT,invbanktran)

def flatbank(g,f,acctid,context,cc=False):
    if cc:
        aCCTTYPE="CC"
        invstmtrT=["CCACCTFROM:ACCTID","CURDEF","LEDGERBAL:DTASOF"]
    else:
        aCCTTYPE="BANK"
        invstmtrT=["BANKACCTFROM:ACCTID","+BANKID","+ACCTTYPE","CURDEF",
                   "LEDGERBAL:DTASOF"]
    invstmtrs=g.query("select "+vars(invstmtrT)+"{"+pairs("x",invstmtrT)+"}",
                      initNs=NS,initBindings={"ACCTID":acctid})
    assert len(invstmtrs)==1
    for invstmtr in invstmtrs:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        f.add((bNode,A3["aCCTTYPE"],Literal(aCCTTYPE)))
    #BAL
    balT=["LEDGERBAL:BALAMT","+dtASOF","?AVAILBAL:bALAMT","+dTASOF"]
    bals=g.query("select "+vars(balT)+"{"+
                    pairs("x",[aCCTTYPE+"ACCTFROM:ACCTID"])+
                    pairs("x",balT)+"}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for bal in bals:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,balT,bal)
    #BANKTRANLIST
    invtranlistT=["DTSTART","DTEND"]
    invtranlists=g.query("select "+vars(invtranlistT)+"{"+
                    pairs("x",[aCCTTYPE+"ACCTFROM:ACCTID"])+
                    "?x ofx:BANKTRANLIST ?y . "+
                    pairs("y",invtranlistT)+"}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invtranlist in invtranlists: # At most one INVTRANLIST per account
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,invtranlistT,invtranlist)
    
    #STMTTRN
    stmttrnT=["TRNTYPE","DTPOSTED","TRNAMT","FITID",
              "?DTUSER","?DTAVAIL",
              "?CORRECTFITID","?CORRECTACTION","?SRVRTID",
              "?CHECKNUM","?REFNUM","?SIC","?PAYEEID","?NAME",
              "?PAYEE:NAME","+ADDR1","+ADDR2","+ADDR3","+CITY","+STATE",
              "+POSTALCODE","+COUNTRY","+PHONE",
              "?BANKACCTTO:bankid","+acctid","+accttype",
              "?CCACCTTO:acctid",
              "?mEMO","?INV401KSOURCE",
              "?CURRENCY:currate","+cursym",
              "?ORIGCURRENCY:currate","+cursym",
              ]
    invbanktrans=g.query("select "+vars(stmttrnT)+"{"+
                    pairs("x",[aCCTTYPE+"ACCTFROM:ACCTID"])+
                    "?x ofx:BANKTRANLIST ?y ."+
                    "?y ofx:STMTTRN ?z ."+
                    pairs("z",stmttrnT)+
                    "}",
                    initNs=NS,initBindings={"ACCTID":acctid})
    for invbanktran in invbanktrans:
        bNode=BNode()
        f.add((context,flatLine,bNode))
        bprt(f,bNode,invstmtrT,invstmtr)
        bprt(f,bNode,stmttrnT,invbanktran)

def flat(fin,fout,context=None,hashing=False):    
    g = Graph()
    g.parse(fin,format='n3',publicID=context)
    f = Graph() #output graph
    f.bind("ofx",OFX)
    f.bind("xsd",XSD)
    f.bind("a3",A3)
    
    context=FS[context]
    flatseclist(g,f,context)
    
    accounts=g.query("select ?typ "+vars(["ACCTID"])+"{"+
            "[ ?typ ?x ] . "+
            pairs("x",["ACCTID"])+"}",initNs=NS)    
    for account in accounts:
        acctid=account[1]
        typ=account[0]
        if typ==OFX["BANKACCTFROM"]:
            logging.info("Bank statement")
            flatbank(g,f,acctid,context)
        elif typ==OFX["CCACCTFROM"]:
            logging.info("Credit card statement")
            flatbank(g,f,acctid,context,cc=True)
        elif typ==OFX["INVACCTFROM"]:
            logging.info("Investment statement")
            flatinv(g,f,acctid,context)
        else:
            logging.info("Unrecognized type %s"%account[1])
    
    fout.write(f.serialize(format='n3'))
